package leetcode_2401_2500;

import node.TreeNode;

import java.util.*;

public class LeeCode_2458 {
    public static void main(String[] args) {
//        System.out.println(Arrays.toString(treeQueries(new TreeNode(1, new TreeNode(2), new TreeNode(3)), new int[]{2, 3})));
        System.out.println(Arrays.toString(treeQueries(new TreeNode(2, new TreeNode(1), new TreeNode(4, new TreeNode(3), null)),
                new int[]{1, 4, 3, 4})));
//        System.out.println(Arrays.toString(treeQueries(new TreeNode(1, null, new TreeNode(5, new TreeNode(3, new TreeNode(2), new TreeNode(4)), null)),
//                new int[]{3, 5, 4, 2, 4})));
    }

    private static int[] treeQueries(TreeNode root, int[] queries) {
        Set<Integer> most = new HashSet<>();
        Map<Integer, Integer> map = new HashMap<>();
        Map<Integer, Integer> cs = new HashMap<>();
        Map<Integer, Set<Integer>> tm = new HashMap<>();
        dfs(root, cs, tm, map, 0);
        get(root, most, cs);
        for (int i = 0; i < queries.length; i++) {
            int q = queries[i];
            if (!most.contains(q)) {
                queries[i] = cs.get(root.val);
            } else {
                int d = map.get(q), res = -1;
                for (Integer node : tm.get(d)) {
                    if (node == q)
                        continue;
                    res = Math.max(res, cs.get(node));
                }
                queries[i] = d + res;
            }
        }
        return queries;
    }

    private static void get(TreeNode node, Set<Integer> most, Map<Integer, Integer> cs) {
        most.add(node.val);
        int l = -1, r = -1;
        if (node.left != null) {
            l = cs.get(node.left.val);
        }
        if (node.right != null) {
            r = cs.get(node.right.val);
        }
        if (l == r)
            return;
        if (l > r)
            get(node.left, most, cs);
        else
            get(node.right, most, cs);
    }

    private static int dfs(TreeNode node, Map<Integer, Integer> cs, Map<Integer, Set<Integer>> tm, Map<Integer, Integer> map, int d) {
        if (node == null)
            return -1;
        int cl = dfs(node.left, cs, tm, map, d + 1);
        int cr = dfs(node.right, cs, tm, map, d + 1);
        int val = node.val;
        cs.put(val, Math.max(cl, cr) + 1);
        if (tm.containsKey(d)) {
            tm.get(d).add(val);
        } else {
            Set<Integer> set = new HashSet<>();
            set.add(val);
            tm.put(d, set);
        }
        map.put(val, d);
        return cs.get(val);
    }
}